package edu.indiana.d2i.cluster;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.UnsupportedEncodingException;

import com.jcraft.jsch.Channel;
import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.ChannelSftp;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.jcraft.jsch.SftpException;
import com.jcraft.jsch.SftpProgressMonitor;

public class SSHSession {
	private JSch jsch = new JSch();
	private Session session;
	private StringBuilder cmdPool = new StringBuilder();

//	private final String SUDO_PREFIX = "sudo -S -p '' ";
	private final String SUDO_PREFIX = "sudo ";

	private int getCmdChannelScreenOutput(ChannelExec channel)
			throws IOException {
		InputStream in = channel.getInputStream();
		// don't disconnect too early!
		byte[] tmp = new byte[1024];
		while (true) {
			while (in.available() > 0) {
				int i = in.read(tmp, 0, 1024);
				if (i < 0)
					break;
				System.out.print(new String(tmp, 0, i));
			}
			if (channel.isClosed()) {
				System.out.println("exit-status: " + channel.getExitStatus());
				return channel.getExitStatus();
			}
			try {
				Thread.sleep(1000);
			} catch (Exception ee) {
			}
		}
	}

	public SSHSession(String host, String user, String passwd)
			throws JSchException {
		java.util.Properties config = new java.util.Properties();
		config.put("StrictHostKeyChecking", "no");
		session = jsch.getSession(user, host, 22);
		session.setPassword(passwd);
		session.setConfig(config);
		session.connect();
	}

	public void close() throws Exception {
		flush();
		session.disconnect();
	}

	public int flush() throws Exception {
		int exitcode = 0;
		if (cmdPool.length() > 0) {
			Channel channel = session.openChannel("exec");
			((ChannelExec) channel).setCommand(cmdPool.toString());
			((ChannelExec) channel).setErrStream(System.err);

			channel.connect();
			exitcode = getCmdChannelScreenOutput((ChannelExec) channel);
			channel.disconnect();
		}
		cmdPool.delete(0, cmdPool.length());
		return exitcode;
	}

	public void writeSmallFile(String content, String pathname, boolean sudo)
			throws Exception {
		writeFile(new ByteArrayInputStream(content.getBytes("UTF-8")), pathname, sudo);
	}

	public void writeFile(InputStream input, String pathname, boolean sudo) throws Exception {
		ChannelSftp sftpchannel = (ChannelSftp) session.openChannel("sftp");
		sftpchannel.connect();
		if (sudo) {
			Channel cmdchannel = session.openChannel("exec");
			((ChannelExec) cmdchannel).setCommand(SUDO_PREFIX + "rm -fr "
					+ pathname + ";" + SUDO_PREFIX + "touch " + pathname + ";"
					+ SUDO_PREFIX + "chmod o+w " + pathname + ";" + "exit;");
			cmdchannel.connect();
			getCmdChannelScreenOutput((ChannelExec) cmdchannel);
			cmdchannel.disconnect();
		}
		sftpchannel.put(input, pathname, ChannelSftp.OVERWRITE);
		sftpchannel.disconnect();
	}
	
	public void execCmd(String cmd, boolean sudo) {
		if (cmdPool.length() > 0) cmdPool.append(";");
		if (sudo)
			cmdPool.append(SUDO_PREFIX + cmd);
		else
			cmdPool.append(cmd);
	}
	
	public void execCmdWithoutBlocking(String cmd, boolean sudo) throws Exception {
		String command = "";
		if (sudo)
			command = SUDO_PREFIX + cmd;
		else
			command = cmd;
		Channel channel = session.openChannel("exec");
		((ChannelExec) channel).setCommand(command);
		Thread.sleep(2000);

		channel.connect();
		channel.disconnect();
	}
	
	public StringBuffer execCmdReturnScreenOutput(String cmd, boolean sudo) throws Exception {
		StringBuffer buf = new StringBuffer();
		
		Channel channel = session.openChannel("exec");
		((ChannelExec) channel).setCommand(cmd);
		((ChannelExec) channel).setErrStream(System.err);

		channel.connect();
		InputStream in = channel.getInputStream();
		// don't disconnect too early!
		byte[] tmp = new byte[1024];
		while (true) {
			while (in.available() > 0) {
				int i = in.read(tmp, 0, 1024);
				if (i < 0) break;
				buf.append(new String(tmp, 0, i));
			}
			if (channel.isClosed()) {
				System.out.println("exit-status: " + channel.getExitStatus());
				break;
			}
			Thread.sleep(1000);
		}
		channel.disconnect();
		
		return buf;
	}
}
